#include <stdio.h>
#include <Windows.h>
#include <Shlwapi.h>

#include "builder.h"
#include "exec.h"
#include "getopt.h"
#include "helper.h"
#include "rc4.h"

#define DEFAULT_NAME "out.exe"

#pragma comment(lib, "Shlwapi.lib")

void PrintHelp(LPSTR self) {
	PathStripPath(self);

	printf(
		"Usage:\n"
		"\t%s [-p <PAYLOAD PATH> [-o <OUT FILE NAME]]\n\n"
		"Example:\n"
		"\t%s -p mypayload.exe\n"
		"\t%s -p mypayload.exe -o loader.exe\n\n",
		self,
		self,
		self
	);
}

/*
 * Usage %s [-p <EXE PAYLOAD PATH> [-o <OUT FILE NAME>]]
 */
int main(int argc, char *argv[]) {
	if (argc == 1) {
		// Do the magics.
		if (!Execute()) {
			PRINT_ERROR("The program terminated unsuccessfully.\n");
			return 1;
		}

		PRINT_SUCCESS("The program terminated successfully.\n");
		return 0;
	} else {
		if (argc >= 3 || argc >= 5) {
			BOOL bRet = FALSE;

			int c = 0;
			LPCSTR szPayloadPath = NULL;
			LPCSTR szNewFileName = DEFAULT_NAME;

			while ((c = getopt(argc, argv, "p:o:")) != -1) {
				switch (c) {
					case 'h':
						PrintHelp(argv[0]);
						return 0;
					case 'p':
						szPayloadPath = optarg;
						break;
					case 'o':
						szNewFileName = optarg;
						break;
					case '?':
						PRINT_WARNING("Unknown option: \"-%c\"\n", optopt);
						break;
					default:
						break;
				}
			}

			// Check mandatory payload argument.
			if (argc < 5 && szPayloadPath == NULL) {
				PRINT_ERROR("Payload argument must exist.\n");
				return 1;
			}

			PRINT_INFO("Building payload...\n");
			PRINT_INFO(
				"Out file name: \"%s\"\n",
				szNewFileName
			);
			PRINT_INFO("Payload: \"%s\"\n",
				szPayloadPath
			);

			bRet = BuildLoader(szNewFileName, szPayloadPath);

			if (!bRet) {
				PRINT_ERROR("The program terminated unsuccessfully.\n");
				return 1;
			} else {
				PRINT_SUCCESS("The program terminated successfully.\n");
				return 0;
			}
		}
	}

	PrintHelp(argv[0]);

	return 0;
}